package informed.sgd;

import java.util.ArrayList;
import java.util.Collections;
import java.util.Random;

import informed.samples.MatrixCellValue;
import informed.samples.Sample;
import cern.colt.matrix.DoubleMatrix2D;
import cern.colt.matrix.impl.DenseDoubleMatrix2D;
import cern.colt.matrix.linalg.Algebra;

public class RegularizedSVD {
	
	private DoubleMatrix2D P;
	private DoubleMatrix2D Q;
	private DoubleMatrix2D cache;
	private double lambda;
	private double learningRate;
	private int factors;
	
	private double error;
	private int iter;
	private int nSamples;
	
	public RegularizedSVD(int n, int m, int factors){
		this.factors = factors;
		iter = 0;
		error = 0;
		P = new DenseDoubleMatrix2D(n,factors);
		Q = new DenseDoubleMatrix2D(factors,m);
		cache = new DenseDoubleMatrix2D(n,m);
	}
	
	private void init(ArrayList<Sample> samples){
		Random r = new Random(System.currentTimeMillis());
		double scale = 0.2;
		for(int f = 0; f < factors; f++){
			for(int n = 0; n < P.rows(); n++){
				P.setQuick(n, f, scale*r.nextDouble());
			}
			for(int m = 0; m < Q.columns(); m++){
				Q.setQuick(f, m, scale*r.nextDouble());
			}
		}
		learningRate = 0.1;
		lambda = 0.0001;
		System.out.println("Initializing P and Q. Learning rate = " + learningRate);
		nSamples = samples.size();
	}
	
	public void run(ArrayList<Sample> samples, int epochs){
		init(samples);
		for(int i = 0; i < factors; i++){
			for(int j = 0; j < epochs; j++){
				Random rand = new Random(System.currentTimeMillis());
				Collections.shuffle(samples, rand);
				for(int k = 0; k < samples.size(); k++){
					move(samples.get(k), i);
					iter++;
				}
			}
			for(Sample x: samples){
				MatrixCellValue s = (MatrixCellValue)x;
				cache.setQuick(s.i, s.j, (cache.getQuick(s.i, s.j) + P.getQuick(s.i, i)*Q.getQuick(i, s.j)));
			}
			error *= 1.0/nSamples;
			System.out.println("Iteration " + iter + " error=" + error);
			error = 0;
		}
	}
	
	private void move(Sample x, int factor){
		MatrixCellValue s = (MatrixCellValue)x;
		double gamma = learningRate;
		
		double e = s.x - ( cache.getQuick(s.i, s.j) + P.getQuick(s.i, factor)*Q.getQuick(factor, s.j) );
		error += e*e;
		
		double gP = gradientP(s,factor,e);
		double gQ = gradientQ(s,factor,e);
		P.setQuick(s.i, factor, P.getQuick(s.i, factor)-gamma*gP);
		Q.setQuick(factor, s.j, Q.getQuick(factor, s.j)-gamma*gQ);
	}
	
	private double gradientP(MatrixCellValue s, int f, double e){
		double eq = -e*Q.getQuick(f, s.j);
		double lp = lambda*P.getQuick(s.i, f);
		return eq+lp;
	}
	
	private double gradientQ(MatrixCellValue s, int f, double e){
		double ep = -e*P.get(s.i, f);
		double lq = lambda*Q.getQuick(f, s.j);
		return ep+lq;
	}
	
	public String printSolution(){
		Algebra algebra = new Algebra();
		String s = "";
		s = "\nMatrix P: ";
		s += P.toString();
		s += "\nMatrix Q: ";
		s += Q.toString();
		s += "\nReconstruction: ";
		s += algebra.mult(P, Q).toString();
		return s;
	}

}
